import glm
import OpenGL.GL as gl
import numpy as np
import pygl

class ParametricSurface:
    def __init__(self, slices = 32, stacks = 32) -> None:
        number_vertex = (slices + 1) * (stacks + 1)
        self.number_indices = slices * stacks * 6
        self.vertices = np.empty((number_vertex, 8), dtype=gl.GLfloat)
        self.indices = np.empty((self.number_indices), dtype=gl.GLuint)
        EPS = 0.00001
        self.slice_count = slices + 1
        for i in range(stacks + 1):
            v = i / float(stacks)
            for j in range(self.slice_count):
                u = j / float(slices)
                index = j + (slices + 1) * i
                x, y, z = self.evaluate(u, v)
                p0 = glm.vec3(x, y, z)
                # 位置
                self.vertices[index, 0] = x
                self.vertices[index, 1] = y
                self.vertices[index, 2] = z
                # 计算法线向量
                if u - EPS >= 0.0:
                    p1 = glm.vec3(self.evaluate(u - EPS, v))
                    pu = p0 - p1
                else:
                    p1 = glm.vec3(self.evaluate(u + EPS, v))
                    pu = p1 - p0
                if v - EPS >= 0:
                    p1 = glm.vec3(self.evaluate(u, v - EPS))
                    pv = p0 - p1
                else:
                    p1 = glm.vec3(self.evaluate(u, v + EPS))
                    pv = p1 - p0                    
                normal = glm.normalize(glm.cross(pu, pv))
                # 法线向量
                self.vertices[index, 3] = normal.x
                self.vertices[index, 4] = normal.y
                self.vertices[index, 5] = normal.z
                # 纹理坐标
                self.vertices[index, 6] = v
                self.vertices[index, 7] = u
        # 生成顶点索引
        index = 0
        for i in range(stacks):
            for j in range(slices):
                a = i * self.slice_count + j
                b = i * self.slice_count + j + 1
                c = (i + 1) * self.slice_count + j + 1
                d = (i + 1) * self.slice_count + j
                self.indices[index + 0] = a
                self.indices[index + 1] = b
                self.indices[index + 2] = d
                self.indices[index + 3] = b
                self.indices[index + 4] = c
                self.indices[index + 5] = d
                index += 6
        self.vao = pygl.VertexArrayObject()
        self.vbo = pygl.VertexBufferObject(self.vertices)
        self.ebo = pygl.ElementBufferObject(self.indices)
        attribute_position = pygl.VertexAttribute("position", 0, 3, 
                                    gl.GL_FLOAT, gl.GL_FALSE, 0)
        attribute_normal = pygl.VertexAttribute("normal", 1, 3, 
            gl.GL_FLOAT, gl.GL_FALSE, 3 * gl.sizeof(gl.GLfloat))
        attribute_tex_coord = pygl.VertexAttribute("tex_coord", 2, 2, 
                gl.GL_FLOAT, gl.GL_FALSE, 6 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexBuffer(self.vbo, 0, 0, 8 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexAttribute(0, attribute_position)
        self.vao.setVertexAttribute(0, attribute_normal)
        self.vao.setVertexAttribute(0, attribute_tex_coord)
        self.vao.setElementBuffer(self.ebo)
        self.rotation = glm.vec3(0.0, 0.0, 0.0)
        self.position = glm.vec3(0.0, 0.0, 0.0)
        self.scale = glm.vec3(1.0, 1.0, 1.0)
        self.updateModelMatrix()

    def render(self)->None:
        self.vao.bind()
        gl.glDrawElements(gl.GL_TRIANGLES, self.number_indices, 
                          gl.GL_UNSIGNED_INT, None)

    def delete(self)->None:
        self.vao.delete()
        self.vbo.delete()
        self.ebo.delete()

    def updateModelMatrix(self)->None:
        scale_matrix = glm.scale(glm.mat4(1.0), self.scale)
        rotation_matrix = glm.rotate(glm.mat4(1.0), glm.radians(self.rotation.x), 
                                     glm.vec3(1.0, 0.0, 0.0))
        rotation_matrix = glm.rotate(rotation_matrix, 
                        glm.radians(self.rotation.y), glm.vec3(0.0, 1.0, 0.0))
        rotation_matrix = glm.rotate(rotation_matrix, 
                        glm.radians(self.rotation.z), glm.vec3(0.0, 0.0, 1.0))
        translation_matrix = glm.translate(glm.mat4(1.0), self.position)
        self.model_matrix = translation_matrix * rotation_matrix * scale_matrix

    def evaluate(self, u: float, v: float) -> tuple:
        pass

class AppleSurface(ParametricSurface):
    def evaluate(self, u: float, v: float) -> tuple:
        u = 2 * np.pi * u
        v = 2 * np.pi * (v - 0.5)
        x = np.cos(u) * (4 + 3.8 * np.cos(v))
        y = np.sin(u) * (4 + 3.8 * np.cos(v))
        z = (np.cos(v) + np.sin(v) - 1) * (1 + np.sin(v)) * \
            np.log(1 - np.pi * v / 10) + 7.5 * np.sin(v)
        return x, y, z     

class KleinSurface(ParametricSurface):
    def evaluate(self, u: float, v: float) -> tuple:
        u *= 2 * np.pi
        v *= 2 * np.pi
        r = 2 * (1 - np.cos(u)/2.0)
        a = 3.0
        b = -8.0
        if u <= np.pi:
            x = a * np.cos(u) * (1 + np.sin(u)) + r * np.cos(u) * np.cos(v)
            y = b * np.sin(u) - r * np.sin(u) * np.cos(v)
        else:
            x = a * np.cos(u) * (1 + np.sin(u)) + r * np.cos(v + np.pi) 
            y = b * np.sin(u)
        z = r * np.sin(v)
        return x, y, z
        
class BreatherSurface(ParametricSurface):
    def evaluate(self, u: float, v: float) -> tuple:
        u = 13.2 * 2 * (u - 0.5)
        v = 37.4 * 2 * (v - 0.5)
        a = 0.4
        r = 1.0 - a**2
        w = np.sqrt(r)
        d = a * ((w * np.cosh(a * u))**2 + (a * np.sin(w * v))**2)
        x = -u + (2.0 * r * np.cosh(a * u) * np.sinh(a * u) / d) 
        y = 2.0 * w * np.cosh(a * u) * (-(w * np.cos(v) * np.cos(w * v)) - 
                                        np.sin(v) * np.sin(w * v)) / d
        z = 2.0 * w * np.cosh(a * u) * (-(w * np.sin(v) * np.cos(w * v)) + 
                                        np.cos(v) * np.sin(w * v)) / d
        return x, y, z







